#ifndef PCPS_WAVEFUNCTION_HEADER
#define PCPS_WAVEFUNCTION_HEADER

#include <src/pcps.h>
#include <src/correlator.h>

namespace pcps {
  
  //-------------------------------------------------------------------------------
  // pcps::Wavefunction -- class to store and access wavefunction data
  //-------------------------------------------------------------------------------

  template <class S> class Wavefunction {

    // constructors
    public:
      Wavefunction(const pcps::Input &userinp, const pcps::Hamiltonian<S> &ham);

    // public member functions
    public:
      std::vector< pcps::Correlator<S> > & correlators() { return _correlators; }
      const std::vector< pcps::Correlator<S> > & correlators() const { return _correlators; }
      std::vector< pcps::Correlator<S> > & pert_corrs() { return _pert_corrs; }
      const std::vector< pcps::Correlator<S> > & pert_corrs() const { return _pert_corrs; }
      std::vector<S> & hf_rdm() { return _hf_rdm; }
      const std::vector<S> & hf_rdm() const { return _hf_rdm; }
      std::vector<S> & restricted_orb_coeffs() { return _restricted_orb_coeffs; }
      const std::vector<S> & restricted_orb_coeffs() const { return _restricted_orb_coeffs; }
      std::vector<S> & agp_weights() { return _agp_weights; }
      const std::vector<S> & agp_weights() const { return _agp_weights; }
      std::vector<S> & pairing_matrix() { return _pairing_matrix; }
      const std::vector<S> & pairing_matrix() const { return _pairing_matrix; }

      // funciton to compute the number of variables in the wavefunction
      size_t nvars() const {
        size_t retval = 0;
        for (typename std::vector< pcps::Correlator<S> >::const_iterator corr = _correlators.begin(); corr != _correlators.end(); corr++)
          if (!corr->fixed() && !corr->ti())
            retval += size_t(corr->nelements());
        return retval;
      }

      // function to get the index of a translationally invariant correlator's parent
      int get_ti_parent_index(const std::vector<int> & sites) const {
        std::vector<int> matches;
        for (int i = 0; i < _correlators.size(); i++)
          if (!_correlators.at(i).ti() && _correlators.at(i).site_vec() == sites)
            matches.push_back(i);
        if (matches.size() != 1)
          throw pcps::Exception( ( boost::format("%i parents were found in pcps::Wavefunction::get_ti_parent_index. There should be exactly 1 parent.")
                                   % matches.size() ).str() );
        return matches.at(0);
      }

      // function to return the positions of each correlator's variables in the ordering of all wavefunction variables
      const std::vector<int> & corr_offsets() const { return _corr_offsets; }

      // function to symmetrize the pairing matrix
      void symmetrize_pairing_matrix() {
        // get the dimension of the pairing matrix
        int n = 0;
        while ( n*n != _pairing_matrix.size() && n < _pairing_matrix.size() ) n++;
        if (n*n != _pairing_matrix.size())
          throw pcps::Exception( (boost::format("pairing matrix not square in pcps::Wavefunction::symmetrize_pairing_matrix")).str() );
        // symmetrize the pairing matrix
        for (int i = 0; i < n; i++)
          for (int j = 0; j < i; j++) {
            const S x = 0.5 * ( _pairing_matrix[i*n+j] + _pairing_matrix[j*n+i] );
            _pairing_matrix[i*n+j] = x;
            _pairing_matrix[j*n+i] = x;
          }
      }

      // functions to return vectors of correlators touching certain sites
      const std::vector< pcps::Correlator<S> *> & touching_correlators(int site0) const {
        return *_1site_touching_lists[site0];
      }
      const std::vector< pcps::Correlator<S> *> & touching_correlators(int site0, int site1) const {
        if (site0 > site1)
          std::swap(site0, site1);
        return *_2site_touching_lists[site0 * _nsites + site1];
      }
      const std::vector< pcps::Correlator<S> *> & touching_correlators(int site0, int site1, int site2, int site3) const {
        throw pcps::Exception( (boost::format("4-site version of pcps::Wavefunction::touching_correlators not yet implemented")).str() );
        int sites[4];
        sites[0] = site0;
        sites[1] = site1;
        sites[2] = site2;
        sites[3] = site3;
        std::sort(&sites[0], &sites[0] + 4);
        return *_4site_touching_lists[_nsites*(_nsites*(_nsites*sites[0] + sites[1]) + sites[2]) + sites[3]];
      }

      // function to initialize the lists of which correlators touch which sites
      void init_touch_vecs();

      // functions to write/read the correlators to/from a binary file
      void dump_binary(const pcps::Input & userinp, const std::string & fname = "pcps_wfn.bin") const;
      //void read_binary(const pcps::Input & userinp, const std::string & filename);

      // function to dump the correlators to a file for use with Jesse's code
      //void dump_for_jesse(const pcps::Input & userinp, const std::string & filename) const;

      // functions to convert the correlator variables to/from a vector
      void to_vector(const pcps::Input & userinp, double * const v) const;
      void from_vector(const pcps::Input & userinp, const double * const v);

      // function to scale the wavefunction's variables
      void scale_variables();

      // function to print the wavefunction
      void print(const pcps::Input & userinp);

    // private member functions
    private:

      // functions used in wavefunction initialization
      void get_text_correlators(const pcps::Input & userinp);
      void get_binary_correlators(const pcps::Input & userinp);
      void get_classical_correlators(const pcps::Input & userinp);
      void get_lattice_correlators(const pcps::Input & userinp);
      void get_all_pair_correlators(const pcps::Input & userinp, std::vector<pcps::Correlator<S> > & correlators);
      void init_correlators(const pcps::Input & userinp);
      void init_orb_coeff(const pcps::Input & userinp, const pcps::Hamiltonian<S> & ham);
      void init_agp_weights(const pcps::Input & userinp);
      void init_pairing_matrix(const pcps::Input & userinp, const pcps::Hamiltonian<S> & ham);

    // private data
    private:

      // number of sites in the lattice
      const int _nsites;

      // correlators
      std::vector< pcps::Correlator<S> > _correlators;

      // perturbative correlators
      std::vector< pcps::Correlator<S> > _pert_corrs;

      // 1-body RDM of the Hartree-Fock reference function
      std::vector<S> _hf_rdm;

      // restricted orbital coefficients
      std::vector<S> _restricted_orb_coeffs;

      // agp orbital weights
      std::vector<S> _agp_weights;

      // pairing matrix
      std::vector<S> _pairing_matrix;

      // positions of each correlator's variables in the ordering of all wavefunction variables
      std::vector<int> _corr_offsets;

      // lists of which correlators touch which sites
      std::vector< boost::shared_ptr< std::vector< pcps::Correlator<S> *> > > _1site_touching_lists;
      std::vector< boost::shared_ptr< std::vector< pcps::Correlator<S> *> > > _2site_touching_lists;
      std::vector< boost::shared_ptr< std::vector< pcps::Correlator<S> *> > > _4site_touching_lists;

    // disabled functions
    private:
      Wavefunction(); // disable default constructor
      Wavefunction(const Wavefunction<S> &); // disable copy constructor
      Wavefunction<S> &operator=(const Wavefunction<S> &); // disable assignment operator
  };

}

#endif
